/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.openhitls.core.CryptoNative;

import javax.net.ssl.SSLPeerUnverifiedException;
import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSessionContext;
import javax.security.cert.X509Certificate;
import java.security.Principal;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.io.ByteArrayInputStream;
import java.util.HashMap;
import java.util.Map;

/**
 * SSL Session implementation backed by openHiTLS.
 */
public class HiTlsSSLSession implements SSLSession {

    private final HiTlsSSLEngine engine;
    private final long creationTime;
    private long lastAccessedTime;
    private String cipherSuite;
    private String protocol;
    private byte[] sessionId;
    private java.security.cert.Certificate[] peerCertificates;
    private Principal peerPrincipal;
    private Principal localPrincipal;
    private final Map<String, Object> values = new HashMap<>();
    private boolean valid = true;

    HiTlsSSLSession(HiTlsSSLEngine engine) {
        this.engine = engine;
        this.creationTime = System.currentTimeMillis();
        this.lastAccessedTime = creationTime;
    }

    void updateFromNative(long sslPtr) {
        if (sslPtr == 0) return;

        // Get protocol version
        int version = CryptoNative.sslGetVersion(sslPtr);
        if (version == CryptoNative.TLS_VERSION_1_3) {
            protocol = "TLSv1.3";
        } else if (version == CryptoNative.TLS_VERSION_1_2) {
            protocol = "TLSv1.2";
        } else {
            protocol = "TLS";
        }

        // Get cipher suite
        cipherSuite = CryptoNative.sslGetCipherSuite(sslPtr);
        if (cipherSuite == null) {
            cipherSuite = "TLS_NULL_WITH_NULL_NULL";
        }

        // Get peer certificates
        try {
            peerCertificates = engine.getPeerCertificates();
            if (peerCertificates != null && peerCertificates.length > 0) {
                java.security.cert.X509Certificate cert =
                    (java.security.cert.X509Certificate) peerCertificates[0];
                peerPrincipal = cert.getSubjectX500Principal();
            }
        } catch (Exception e) {
            // Peer certificates may not be available
        }

        lastAccessedTime = System.currentTimeMillis();
    }

    @Override
    public byte[] getId() {
        return sessionId != null ? sessionId.clone() : new byte[0];
    }

    @Override
    public SSLSessionContext getSessionContext() {
        return null;
    }

    @Override
    public long getCreationTime() {
        return creationTime;
    }

    @Override
    public long getLastAccessedTime() {
        return lastAccessedTime;
    }

    @Override
    public void invalidate() {
        valid = false;
    }

    @Override
    public boolean isValid() {
        return valid;
    }

    @Override
    public void putValue(String name, Object value) {
        if (name == null || value == null) {
            throw new IllegalArgumentException("Name and value cannot be null");
        }
        values.put(name, value);
    }

    @Override
    public Object getValue(String name) {
        if (name == null) {
            throw new IllegalArgumentException("Name cannot be null");
        }
        return values.get(name);
    }

    @Override
    public void removeValue(String name) {
        if (name == null) {
            throw new IllegalArgumentException("Name cannot be null");
        }
        values.remove(name);
    }

    @Override
    public String[] getValueNames() {
        return values.keySet().toArray(new String[0]);
    }

    @Override
    public Certificate[] getPeerCertificates() throws SSLPeerUnverifiedException {
        if (peerCertificates == null || peerCertificates.length == 0) {
            throw new SSLPeerUnverifiedException("Peer not verified");
        }
        return peerCertificates.clone();
    }

    @Override
    public Certificate[] getLocalCertificates() {
        // Return local certificates if configured
        return null;
    }

    @Override
    @SuppressWarnings("deprecation")
    public X509Certificate[] getPeerCertificateChain() throws SSLPeerUnverifiedException {
        Certificate[] certs = getPeerCertificates();
        X509Certificate[] x509Certs = new X509Certificate[certs.length];

        try {
            for (int i = 0; i < certs.length; i++) {
                byte[] encoded = certs[i].getEncoded();
                x509Certs[i] = X509Certificate.getInstance(encoded);
            }
            return x509Certs;
        } catch (Exception e) {
            throw new SSLPeerUnverifiedException("Failed to get certificate chain: " + e.getMessage());
        }
    }

    @Override
    public Principal getPeerPrincipal() throws SSLPeerUnverifiedException {
        if (peerPrincipal == null) {
            throw new SSLPeerUnverifiedException("Peer not verified");
        }
        return peerPrincipal;
    }

    @Override
    public Principal getLocalPrincipal() {
        return localPrincipal;
    }

    @Override
    public String getCipherSuite() {
        return cipherSuite != null ? cipherSuite : "TLS_NULL_WITH_NULL_NULL";
    }

    @Override
    public String getProtocol() {
        return protocol != null ? protocol : "TLS";
    }

    @Override
    public String getPeerHost() {
        return engine.getPeerHost();
    }

    @Override
    public int getPeerPort() {
        return engine.getPeerPort();
    }

    @Override
    public int getPacketBufferSize() {
        // Maximum TLS record size
        return 16384 + 2048; // Record + overhead
    }

    @Override
    public int getApplicationBufferSize() {
        // Maximum application data per record
        return 16384;
    }
}
